package cn.cew.bytespy;

import cn.cew.bytespy.comm.AgentParam;
import cn.cew.bytespy.comm.Type;
import cn.cew.bytespy.transfer.MethodCostTransformer;
import cn.cew.bytespy.transfer.RecoverAgentTransformer;
import cn.cew.bytespy.utils.BaseUtils;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.File;
import java.io.FileOutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.util.*;

/**
 * @author chenenwei
 * @version 1.0.0
 * @description TODO
 * @createTime 2022-10-13 11:08:00
 */
public class SpyApplication {

    public static void agentmain(String agentArgs, Instrumentation inst) throws Exception {

        AgentParam agentParam = new AgentParam(agentArgs);
        String socketClientName = "ByteSpySocketClient";
        String socketClientPackage = "";

        Class<?>[] loadedClasses = inst.getAllLoadedClasses();
        List<Class<?>> list = new ArrayList<>(1);
        for (Class<?> loadedClass : loadedClasses) {
            if (loadedClass.getSimpleName().equals(socketClientName)) {
                socketClientPackage = BaseUtils.packageToDot(loadedClass.getPackage().getName());
                agentParam.getImportPackages().add(socketClientPackage + "." + socketClientName);
                break;
            }
        }

        // 过滤满足条件的类
        for (Class<?> loadedClass : loadedClasses) {
            if (loadedClass.getSimpleName().equals(agentParam.getClassName())) {
                list.add(loadedClass);
            }
        }

        // 判断是否创建过socket类
        if (BaseUtils.isEmpty(socketClientPackage) && BaseUtils.isNotEmpty(list)) {
            // 创建socket client类
            try {
                socketClientPackage = BaseUtils.packageToDot(list.get(0).getPackage().getName());
                String name = socketClientPackage + "." + socketClientName;
                makeSocketClientClass(name);
                agentParam.getImportPackages().add(name);
            } catch (Exception e) {
                System.out.println("创建socket类失败");
                e.printStackTrace();
                return;
            }
        }

        ClassFileTransformer transformer = null;
        if (agentParam.isRecover()) {
            for (Class<?> clz : list) {
                transformer = new RecoverAgentTransformer();
                inst.addTransformer(transformer, true);
                inst.retransformClasses(clz);
            }
        } else {
            if (agentParam.getType().equals(Type.PRINT_PARAM.getValue())) {
                //方法耗时
                for (Class<?> clz : list) {
                    transformer = new MethodCostTransformer(agentParam);
                    inst.addTransformer(transformer, true);
                    inst.retransformClasses(clz);
                }

            } else if (agentParam.getType().equals(Type.METHOD_RETURN.getValue())) {
                // 方法返回值

            }
        }
    }

    /**
     * 创建客户端socket
     * @param className
     */
    private static void makeSocketClientClass (String className) throws Exception{

        String clientCode = " public static void send (String host, String port, String message) {\n" +
                "        \n" +
                "        try {\n" +
                "            int p = Integer.parseInt(port);\n" +
                "            Socket client = new Socket(host, p);\n" +
                "            client.setSoTimeout(10 * 1000);\n" +
                "            Writer writer = new OutputStreamWriter(client.getOutputStream(), \"UTF-8\");\n" +
                "            writer.write(message);\n" +
                "            writer.write(\"EOF\");\n" +
                "            writer.flush();\n" +
                "            writer.close();\n" +
                "            client.close(); \n" +
                "        } catch (Exception e) {\n" +
                "            System.out.println(\"连接socket服务超时\");\n" +
                "            e.printStackTrace();\n" +
                "        }\n" +
                "    }";

        ClassPool classPool = ClassPool.getDefault();
        CtClass tempClass = classPool.makeClass(className);
        classPool.importPackage("java.io.IOException");
        classPool.importPackage("java.io.OutputStreamWriter");
        classPool.importPackage("java.io.Writer");
        classPool.importPackage("java.net.Socket");

        CtMethod cew = CtMethod.make(clientCode, tempClass);
        tempClass.addMethod(cew);
        tempClass.toBytecode();
        tempClass.toClass().newInstance();
        FileOutputStream out = new FileOutputStream(new File("ByteSpySocketClient.class"));
        out.write(tempClass.toBytecode());
        out.flush();
        out.close();

    }
}
